import shutil
import os
import re

import time
import datetime
import random

import torch

from src import UNet
from train_utils import train_one_epoch, evaluate, create_lr_scheduler,addimages
from my_dataset import DriveDataset
import transforms as T
from compute_mean_std import compute_mean_std
import threading
import argparse
import mainui


class SegmentationPresetTrain:#训练图像预处理方法
    def __init__(self, base_size, crop_size, hflip_prob=0.5, vflip_prob=0.5,#base_size：将图像调整到固定的尺寸，以便在训练过程中保持图像大小的一致性
                 mean=0.36,std=0.18):              #crop_size：在图像分割任务中，通常需要将原始图像裁剪成固定大小的图像块，用作模型的输入。crop_size 指定了裁剪后图像的目标尺寸，确保所有训练样本的输入图像大小一致。
        min_size = int(0.5 * base_size)                                                               #hflip_prob：水平翻转的概率 vflip_prob：垂直翻转的概率
        max_size = int(1.2 * base_size)                                                              #mean(均值)和std(标准差):在图像处理中，通常会计算数据集中所有图像的像素均值和方差，然后将这个均值应用于图像数据的标准化过程中。(original-mean)/std，三个值为在 ImageNet 数据集上计算得出的图像通道（RGB）均值和方差
                                                                                                                         #min_size 和 max_size 是用于控制图像尺寸变换的两个重要参数，通常用于数据预处理和图像增强过程中。
                                                                                                                         # min_size：可以通过将图像的短边（即宽度或高度中较小的那个值）缩放至 min_size 来确保所有图像在进行尺寸调整后具有相同的最小尺寸。这个操作通常用于数据预处理阶段，以便在输入模型之前使所有图像达到相同的最小尺寸。
                                                                                                                         #max_size ：当图像缩放后的任一边（宽度或高度）超过max_size时，会对图像进行进一步的调整，确保图像的最大尺寸不超过指定的max_size。这个操作有助于限制图像的尺寸，以避免模型输入过大而导致内存消耗过高或计算负担增加。
        trans = [T.RandomResize(min_size, max_size)]                                     #将图像的最小边长缩放到min_size和max_size中的一个值的大小
        if hflip_prob > 0:
            trans.append(T.RandomHorizontalFlip(hflip_prob))                           #对image和target进行水平翻转
        if vflip_prob > 0:
            trans.append(T.RandomVerticalFlip(vflip_prob))
        trans.extend([
            T.RandomRotate(angles=[-30, -15, 0, 15, 30]),
            T.RandomCrop(crop_size),
            T.ToTensor(),                                                                                        #用于将 PIL 图像（Pillow 库加载的图像对象）或 ndarray 数组转换为张量（tensor）：1、如果输入是 PIL 图像对象（PIL.Image.Image），则会将图像的通道顺序从 RGB 调整为 PyTorch 预期的顺序（即将通道顺序从 HWC 调整为 CHW，其中 H 表示图像高度，W 表示图像宽度，C 表示通道数）。2、将图像的像素值（0~255 范围内的整数值）标准化到 [0, 1] 的范围内，以便与神经网络模型的输入要求相匹配。
            T.Normalize(mean=mean, std=std),
        ])
        self.transforms = T.Compose(trans)                                                        #trans列表中，每个元素代表一种预处理，compose打包复制给transforms

    def __call__(self, img, target):
        return self.transforms(img, target)


class SegmentationPresetEval:#验证图像预处理方法
    def __init__(self, mean=0.36, std=0.18):
        self.transforms = T.Compose([
            T.ToTensor(),
            T.Normalize(mean=mean, std=std),
        ])

    def __call__(self, img, target):
        return self.transforms(img, target)


def get_transform(train, mean=0.36, std=0.18):
    base_size = 512                   #原图是565*584，通常情况下，base_size 是根据图像的宽度或高度来设定的，以确保图像处理过程中可以保持一定的尺寸比例或满足特定的输入要求。
    crop_size = 480                   #输入的图像为480*480

    if train:
        return SegmentationPresetTrain(base_size, crop_size, mean=mean, std=std)       #如果是在train，就返回对训练集的预处理方法
    else:
        return SegmentationPresetEval(mean=mean, std=std)


def create_model(num_classes):
    model = UNet(in_channels=1, num_classes=num_classes, base_c=64)       #输入通道数为3，分类类别数为2，基础通道数（第一层的卷积操作将使用基础通道数作为卷积核的数量）为64
    return model                                                                                                 #U-Net没有什么预训练权重

#文件操作功能函数
def find_max_numeric_filename(folder_path):                                             #找到training中最大文件名
    max_number = float('-inf')  # 初始化最大数字为负无穷
    # 遍历文件夹中的所有文件
    for filename in os.listdir(folder_path):
        if filename.endswith('.png'):  # 确保是PNG图片文件
            try:
                # 提取文件名中的数字部分并转换为整数
                number = int(os.path.splitext(filename)[0])
                if number > max_number:
                    max_number = number
            except ValueError:
                continue  # 文件名不是数字，跳过
    return  max_number

def copy_and_rename_image(source_folder, target_folder, old_filename,new_filename):   #将redundance中的文件复制到training中并重命名
    # 源文件路径
    source_image_path = os.path.join(source_folder, old_filename)
    # 目标文件路径（新的文件名）
    target_image_path = os.path.join(target_folder, new_filename)
    # 复制文件并重命名
    shutil.copy(source_image_path, target_image_path)

def get_nth_png_filename(folder_path, n):                                                    #拿到第n张图片的文件名，注意这里是1,10,11,12,13,2,20,21...这种顺序的第n张
    try:
        png_files = [file for file in os.listdir(folder_path) if file.endswith('.png')]
        if n > 0 and n <= len(png_files):
            match = re.search(r'\d+', png_files[n-1])
            return int(match.group()) # n-1 因为列表索引从0开始
        else:
            print(f"Invalid value of n. There are {len(png_files)} PNG files in the folder.")
            return None
    except Exception as e:
        print(f"Error: {e}")
        return None

def delete_png_file(file_path):                                                                       #删除redundance中弄去training的图片
    try:
        if os.path.exists(file_path) and file_path.endswith('.png'):
            os.remove(file_path)
            print(f"Moved  {file_path}")
        else:
            print(f"File {file_path} not found or not a PNG file")
    except Exception as e:
        print(f"Error deleting file: {e}")

def main(args,pause_event,update_ui_callback):
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    batch_size = args.batch_size
    # segmentation nun_classes + background
    num_classes = args.num_classes + 1

    # using compute_mean_std.py
    mean,std = compute_mean_std()
    print(f"mean: {mean}")
    print(f"std: {std}")

    # 用来保存训练以及验证过程中信息
    results_file = "results{}.txt".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))   #.format() 方法，可以将变量的值动态地嵌入到字符串中的占位符位置，从而生成所需的格式化输出。

    train_dataset = DriveDataset(args.data_path,
                                 train=True,
                                 rdd=False,
                                 transforms=get_transform(train=True, mean=mean, std=std))

    val_dataset = DriveDataset(args.data_path,
                               train=False,
                               rdd=False,
                               transforms=get_transform(train=False, mean=mean, std=std))
    """
    rdd_dataset = DriveDataset(args.data_path,
                               train=False,
                               rdd=True,
                               transforms=get_transform(train=False, mean=mean, std=std))
   """
    num_workers = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 0])                    #os.cpu_count()：cpu的核数（i7-14700hx-20）[20,4,0]:num_workers不设置为0好像没法运行
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size,
                                               num_workers=num_workers,
                                               shuffle=True,
                                               pin_memory=True,
                                               collate_fn=train_dataset.collate_fn)

    val_loader = torch.utils.data.DataLoader(val_dataset,
                                             batch_size=1,
                                             num_workers=num_workers,
                                             pin_memory=True,
                                             collate_fn=val_dataset.collate_fn)
    """
    rdd_loader = torch.utils.data.DataLoader(rdd_dataset,
                                             batch_size=1,
                                             num_workers=num_workers,
                                             pin_memory=True,
                                             collate_fn=rdd_dataset.collate_fn)
    """

    model = create_model(num_classes=num_classes)
    model.to(device)

    params_to_optimize = [p for p in model.parameters() if p.requires_grad]       #params_to_optimize 列表中存储了模型中需要进行梯度更新的所有参数。在训练过程中，优化器（如 Adam、SGD 等）会针对这些参数计算梯度并更新参数值，从而最小化损失函数并优化模型的性能。这种方式允许灵活地选择需要优化的参数，可以排除不需要更新的参数（如冻结的部分参数或不计算梯度的参数），以提高训练效率和性能。
                                                                                                                                 #model.parameters(): 这是模型对象的方法，用于返回模型中所有的参数。模型的参数通常包括权重矩阵和偏置向量等，在训练过程中需要根据损失函数来更新这些参数。
    optimizer = torch.optim.SGD(                                                                             #p.requires_grad: 这是 PyTorch 中张量对象的属性，表示张量是否需要计算梯度。默认情况下，模型的参数都会设置为需要计算梯度 (requires_grad=True)，以便在反向传播过程中计算梯度并进行参数更新。
        params_to_optimize,
        lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay
    )

    scaler = torch.cuda.amp.GradScaler() if args.amp else None

    # 创建学习率更新策略，这里是每个step更新一次(不是每个epoch)
    lr_scheduler = create_lr_scheduler(optimizer, len(train_loader), args.epochs, warmup=True)

    if args.resume:
        checkpoint = torch.load(args.resume, map_location='cpu')
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        args.start_epoch = checkpoint['epoch'] + 1
        if args.amp:
            scaler.load_state_dict(checkpoint["scaler"])

    best_dice = 0.
    start_time = time.time()
    maxl = 667
    for epoch in range(args.start_epoch, args.epochs):
        pause_event.wait()
        mean_loss, lr = train_one_epoch(model, optimizer, train_loader, device, epoch, num_classes,          #训练一个epoch
                                        lr_scheduler=lr_scheduler, print_freq=args.print_freq, scaler=scaler,pause_event=pause_event)

        add_info = ""
        if epoch>0:                                                                                                                                        #运行15个epoch后(80%)，每epoch进行一次主动学习
            rdd_dataset = DriveDataset(args.data_path,
                                       train=False,
                                       rdd=True,
                                       transforms=get_transform(train=False, mean=mean, std=std))
            rdd_loader = torch.utils.data.DataLoader(rdd_dataset,
                                                     batch_size=1,
                                                     num_workers=num_workers,
                                                     pin_memory=True,
                                                     collate_fn=rdd_dataset.collate_fn)

            max_entropy = addimages(model,rdd_loader,device=device,num_classes=num_classes)        #max_entropy对应的第几张冗余图像是对于dataloder的1,10,11,12,13...这种排列顺序的第23张，不是对应于文件名1,2,3,4,5,6...的第23张

            """
            #随机选择实验
            max_entropy = random.randint(1, maxl)
            maxl -= 1
            """

            folder_path = r'SLIVER07\redundance\images'
            max_entropy_filename = get_nth_png_filename(folder_path, max_entropy)                        #int max_entropy_filename


            add_info = f"#########应该手动标记第{max_entropy_filename}张冗余图像#########"

            #移动图像
            old = max_entropy_filename
            source_folder1 = r'./SLIVER07\redundance\images'
            target_folder1 = r'./SLIVER07\training\images'
            source_folder2 = r'./SLIVER07\redundance\1st_manual'
            target_folder2 = r'./SLIVER07\training\1st_manual'
            source_folder3 = r'./SLIVER07\redundance\mask'
            target_folder3 = r'./SLIVER07\training\mask'
            old_filename = str(old) + '.png'
            #找training\images中文件名中最大的值
            folder_path = r'SLIVER07\training\images'
            max_num = find_max_numeric_filename(folder_path) + 1

            new_filename = str(max_num) + '.png'

            print("#########应该手动标记第", max_entropy_filename, "张冗余图像#########")
            # Pause training and call the UI update callback
            pause_event.clear()
            update_ui_callback(source_folder1, old_filename, target_folder2, new_filename)
            pause_event.wait()  # Wait until the UI signals to continue

            copy_and_rename_image(source_folder1, target_folder1, old_filename, new_filename)
#程序暂停前后端交互：上传标注后的图片到./SLIVER07\training\1st_manual并命名为new_filename，上传对应的mask文件到./SLIVER07\training\mask并命名为new_filename
            #copy_and_rename_image(source_folder2, target_folder2, old_filename, new_filename)
            copy_and_rename_image(source_folder3, target_folder3, old_filename, new_filename)

            # 指定要删除的 PNG 文件路径
            file_to_delete1 =  source_folder1 + "\\"+old_filename
            #file_to_delete2 = source_folder2 + "\\" + old_filename
            file_to_delete3 = source_folder3 + "\\" + old_filename
            # 调用删除函数
            delete_png_file(file_to_delete1)
            #delete_png_file(file_to_delete2)
            delete_png_file(file_to_delete3)

            #别忘了每五个epoch再重新加载training
            if epoch > 18 and (epoch+1)%5 == 0:

                print("#########装入5张手工标记后图像#########")
                add_info = add_info+"\n"+"#########装入5张手工标记后图像#########"

                train_dataset = DriveDataset(args.data_path,
                                            train=True,
                                            rdd=False,
                                            transforms=get_transform(train=True, mean=mean, std=std))

                train_loader = torch.utils.data.DataLoader(train_dataset,
                                                        batch_size=batch_size,
                                                        num_workers=num_workers,
                                                        shuffle=True,
                                                        pin_memory=True,
                                                        collate_fn=train_dataset.collate_fn)

        # 验证
        confmat, dice = evaluate(model, val_loader, device=device, num_classes=num_classes)
        val_info = str(confmat)
        print(val_info)
        print(f"dice coefficient: {dice:.3f}")

        # write into txt
        with open(results_file, "a") as f:
            # 记录每个epoch对应的train_loss、lr以及验证集各指标
            train_info = f"[epoch: {epoch}]\n" \
                         f"train_loss: {mean_loss:.4f}\n" \
                         f"lr: {lr:.6f}\n" \
                         f"dice coefficient: {dice:.3f}\n"
            if epoch>14:
                f.write(train_info + val_info + "\n" + add_info + "\n\n")
            else:
                f.write(train_info + val_info + "\n\n")

        if args.save_best is True:
            if best_dice < dice:
                best_dice = dice
            else:
                continue

        save_file = {"model": model.state_dict(),
                     "optimizer": optimizer.state_dict(),
                     "lr_scheduler": lr_scheduler.state_dict(),
                     "epoch": epoch,
                     "args": args}
        if args.amp:
            save_file["scaler"] = scaler.state_dict()

        if args.save_best is True:
            torch.save(save_file, "save_weights/best_model.pth")
        else:
            torch.save(save_file, "save_weights/model_{}.pth".format(epoch))

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print("training time {}".format(total_time_str))


def parse_args():#将命令行输入的参数解析成 Python 对象，使得我们可以在程序中轻松地访问这些参数值，并根据用户提供的参数执行相应的操作。
    import argparse
    parser = argparse.ArgumentParser(description="pytorch unet training")

    parser.add_argument("--data-path", default="./", help="SLIVER07 root")         #数据集路径
    # exclude background
    parser.add_argument("--num-classes", default=1, type=int)                       #类别为1（main中会加1）：该任务是一个二元分割任务：模型的目标是将图像中的目标区域与背景进行区分，生成一个二值的分割结果。
    parser.add_argument("--device", default="cuda", help="training device")
    parser.add_argument("-b", "--batch-size", default=4, type=int)                #一次迭代用到的样本个数（8G~16-64）
    parser.add_argument("--epochs", default=150, type=int, metavar="N",
                        help="number of total epochs to train")
    #优化器
    parser.add_argument('--lr', default=0.0008, type=float, help='initial learning rate')#初始学习率：用于控制模型参数在每一次迭代（epoch），参数沿着梯度的反方向移动的距离。
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',#Momentum 可以防止梯度更新过程中出现过多的震荡，通过考虑历史梯度的加权平均值，保持参数更新方向的一致性。（0.9/0.95）
                        help='momentum')
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,#代表权重衰减（Weight Decay），也称为 L2 正则化项的系数。在训练神经网络时，通过调整 --wd 参数可以控制模型权重更新过程中的正则化程度，有助于提高模型的泛化能力和抗过拟合能力。
                        metavar='W', help='weight decay (default: 1e-6)',
                        dest='weight_decay')
    parser.add_argument('--print-freq', default=1, type=int, help='print frequency')#指定了多少个批次（batch）训练后，打印一次训练过程中的信息：训练损失（training loss）、准确率（accuracy）、学习率（learning rate）等关键信息
    parser.add_argument('--resume', default='', help='resume from checkpoint')#模型训练将从之前保存的检查点（checkpoint）中恢复，并继续进行训练，而不是从头开始训练模型。
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--save-best', default=True, type=bool, help='only save best dice weights')#用于保存模型的最佳版本或在验证集上表现最好的模型参数。
    # Mixed precision training parameters混合精度训练是一种通过使用较低精度的浮点数（例如16位浮点数）来加快训练速度和减少内存使用的技术，同时保持一定的计算精度以确保模型的准确性。
    parser.add_argument("--amp", default=False, type=bool,                          #启用自动混合精度
                        help="Use torch.cuda.amp for mixed precision training")
    parser.add_argument('--example', type=str, default='example_arg')

    args = parser.parse_args()#解析命令行输入的参数。这将会检查命令行中提供的参数，并根据预先定义的规则将其转换成相应的 Python 对象。

    return args#一旦解析完成，可以通过访问返回的 args 对象来获取用户提供的参数值，并根据需要在程序中进行处理。


if __name__ == '__main__':
    args = parse_args()

    if not os.path.exists("./save_weights"):
        os.mkdir("./save_weights")

    main(args)
